from typing import Literal

from langchain_core.messages import HumanMessage
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import MessagesState, END, StateGraph
from langgraph.prebuilt import ToolNode


@tool
def weather_search(query: str):
    """天气搜索工具"""
    if "上海" in query.lower() or "shanghai" in query.lower():
        return "现在10度，有雾"
    return "现在30度，有雨"


# 工具列表
tools = [weather_search]
# 创建工具列表节点
tool_node = ToolNode(tools=tools)


# 1.初始化模型和工具
def qw_model():
    return ChatOpenAI(
        model='qwen-max-0919',
        base_url='https://dashscope.aliyuncs.com/compatible-mode/v1',
        api_key='sk-debee146f82244268914dd2e3d98761b',
        temperature=0.7,
        top_p=0.9
    )


model = qw_model()
# 将模型绑定工具
qw_model = model.bind_tools(tools)


# 2.定义函数
# 决定是否继续执行（路由函数-条件边的函数）
# Literal定义的两个边
def should_continue(state: MessagesState) -> Literal["tools", END]:
    messages = state["messages"]
    last_message = messages[-1]
    # 如果大模型调用了工具，则转到tool节点，返回标志值
    if last_message.tool_calls:
        return "tools"
    # 否者，停止
    return END


# 定义调用模型的函数
def call_model(state: MessagesState):
    messages = state["messages"]
    resp = qw_model.invoke(messages)
    return {"messages": [resp]}


# 3.初始化图
# 定义状态图
w = StateGraph(MessagesState)
# 4.添加节点
w.add_node("agent", call_model)
w.add_node("tools", tool_node)
# 定义开始节点
w.set_entry_point("agent")

# 5.添加边
# 条件边
w.add_conditional_edges(
    # 首先定义起始节点，这是调用agent节点后采取的
    "agent",
    # 条件边函数
    should_continue,
)

w.add_edge("tools", "agent")

# 6.持久化
memory = MemorySaver()

# 7.编译图
app = w.compile(checkpointer=memory)

final_state = app.invoke(
    {"messages": [HumanMessage(content="你好")]},
    config={"configurable": {"thread_id": "a1"}}
)
res = final_state["messages"][-1].content
print(res)

# final_state_2 = app.invoke(
#     {"messages": [HumanMessage(content="上海天气怎么样")]},
#     config={"configurable": {"thread_id": "a1"}}
# )
# res_2 = final_state_2["messages"][-1].content
# print(res_2)



